from nnunetv2.training.nnUNetTrainer.neuron_seg.NeuronSegBaseTrainer import NeuronSegBaseTrainer, nnUNetTrainerStage4
import torch
from torch import autocast
from nnunetv2.utilities.helpers import empty_cache, dummy_context
from smz.models.unet_sparseconvmoe import UNetWithSparseConvMoE


# Sarse moe
class UNetWithSparseConvMoETrainer(nnUNetTrainerStage4):
    @staticmethod
    def build_network_architecture(
        architecture_class_name, 
        arch_init_kwargs, 
        arch_init_kwargs_req_import, 
        num_input_channels, 
        num_output_channels, 
        enable_deep_supervision = True):

        return UNetWithSparseConvMoE(
            in_channels = num_input_channels,
            out_channels = num_output_channels,
            features = [32, 64, 128, 256],
            spatial_dims = 3,
            num_experts = 6,
            shared_experts = 1,
            top_k = 3,
            gate_hidden = 128,
            tau = 0.5,
            ema_momentum = 0.9,
            eps = 1e-6,
            act = ("swish", {}),
            norm = ("group", {"num_groups": 16, 'affine': True}),
            dropout = None,
            routing_mode = 'hard',
            enable_deep_supervision = enable_deep_supervision,  # Enable nnUNet-style deep supervision
        )   

# dense moe
class UNetWithDenseConvMoETrainer(nnUNetTrainerStage4):
    @staticmethod
    def build_network_architecture(
        architecture_class_name, 
        arch_init_kwargs, 
        arch_init_kwargs_req_import, 
        num_input_channels, 
        num_output_channels, 
        enable_deep_supervision = True):

        return UNetWithSparseConvMoE(
            in_channels = num_input_channels,
            out_channels = num_output_channels,
            features = [32, 64, 128, 256],
            spatial_dims = 3,
            num_experts = 0,
            shared_experts = 4,
            top_k = 1,
            gate_hidden = 128,
            tau = 0.5,
            ema_momentum = 0.9,
            eps = 1e-6,
            act = ("swish", {}),
            norm = ("group", {"num_groups": 16, 'affine': True}),
            dropout = None,
            routing_mode = 'hard',
            enable_deep_supervision = enable_deep_supervision,  # Enable nnUNet-style deep supervision
        )   

# soft moe
class UNetWithSoftConvMoETrainer(nnUNetTrainerStage4):
    @staticmethod
    def build_network_architecture(
        architecture_class_name, 
        arch_init_kwargs, 
        arch_init_kwargs_req_import, 
        num_input_channels, 
        num_output_channels, 
        enable_deep_supervision = True):

        return UNetWithSparseConvMoE(
            in_channels = num_input_channels,
            out_channels = num_output_channels,
            features = [32, 64, 128, 256],
            spatial_dims = 3,
            num_experts = 4,
            shared_experts = 0,
            top_k = 1,
            gate_hidden = 128,
            tau = 0.5,
            ema_momentum = 0.9,
            eps = 1e-6,
            act = ("swish", {}),
            norm = ("group", {"num_groups": 16, 'affine': True}),
            dropout = None,
            routing_mode = 'soft',
            enable_deep_supervision = enable_deep_supervision,  # Enable nnUNet-style deep supervision
        )   